//
// Ephi - simulation of magnetic fields and particles
// Copyright (C) 2007 Indrek Mandre <indrek(at)mare.ee>
// For more information please see http://www.mare.ee/indrek/ephi/
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

#include <assert.h>

#include "dynamics.hpp"
#include "statics.hpp"
#include "consts.hpp"

// i wonder if this is too small, maybe e-12 is better..
#define MY_DT 1e-11

#define ADAPTIVE_STEPS 8
#define ADAPTIVE_ACCF 0.000001

ElectroDynamics::ElectroDynamics (Statics& statics, bool disableThreads) : statics(statics)
{
#ifdef HAVE_THREADS
  if ( disableThreads )
      tmgr = 0;
  else
      tmgr = new TaskManager ();
#endif
  ticks = 0;
  printfp = stdout;
  pmap = 0;
  disable_fields = false;
  setDTM(1.0);

#ifdef HAVE_THREADS
  tasks.resize (THREADS_CPUCOUNT);
  for ( size_t i = 0; i < tasks.size(); i++ )
    {
      tasks[i].n = i;
      tasks[i].dyn = this;
    }
#endif
}

ElectroDynamics::~ElectroDynamics ()
{
  delete [] pmap;
#ifdef HAVE_THREADS
  delete tmgr;
#endif
  for ( size_t i = 0; i < delems.size(); i++ )
      delete delems[i];
}

void ElectroDynamics::setDTM (prec_t accuracy, prec_t stepf)
{
  adaptf = ADAPTIVE_ACCF / accuracy;
  madaptf = 0.2 * adaptf;
  delta_t = MY_DT / stepf;
  update_pmap ();
}

void ElectroDynamics::update_pmap ()
{
  delete [] pmap;
  pmap = new pp[ADAPTIVE_STEPS];
  int c = 0;
  for ( int i = 0; i < ADAPTIVE_STEPS; i++ )
    {
      pp& p = pmap[i];
      if ( !c )
          c = 1;
      else
          c = c * 2;
      p.count = c;
      p.dt = delta_t / p.count;
      p.hdt = delta_t / (2 * p.count);
      p.dt_by_6 = delta_t / (6 * p.count);
      p.dtdt = p.dt * p.dt;
      p.dtdt_by_2 = p.dtdt / 2.0;
      p.dtdt_by_6 = p.dtdt / 6.0;
      p.dtdt_by_8 = p.dtdt / 8.0;
      p.qdt = p.dt / delta_t;
      p.qhdt = p.hdt / delta_t;
//      printf ("%d: %d %g %g\n", i, p.count, p.dt, p.hdt);
    }
}

// q(e)/(4*PI*e0) = q(e)*c^2*m0/(4*PI) = q(e)*c^2*10^-7 = 1.60217653*10^-19 * 299792458^2 * 10^-7
//#define Ke (ELEMENTARY_CHARGE*SPEED_OF_LIGHT*SPEED_OF_LIGHT*MAGNETIC_CONSTANT)
#define Kme (SPEED_OF_LIGHT*SPEED_OF_LIGHT*MAGNETIC_CONSTANT)

void ElectroDynamics::getFields (const vect3d& posv, vect3d& bfield, vect3d& efield, size_t except_electron)
{
  statics.getFields (posv, bfield, efield);

  for ( size_t j = 0; j < delems.size(); j++ )
      delems[j]->addFields (posv, 0, bfield, efield);

  if ( disable_fields )
      return;

  vect3d bf(0, 0, 0), ef(0, 0, 0);
  for ( size_t i = 0; i < particles.size(); i++ )
    {
      if ( i == except_electron )
          continue;
      particle& p = particles[i];
      vect3d r = posv - p.pos;
      prec_t rr = r * r;
      prec_t rlen3 = rr * prec_t_sqrt (rr);
#define HACK 1.0
      prec_t efmod = HACK * p.charge * Kme / rlen3;
      ef += r * efmod;
      prec_t bfmod = (MAGNETIC_CONSTANT * p.charge * HACK) / rlen3;
      bf += bfmod * (p.v % r);
    }
  bfield += bf;
  efield += ef;
}

void ElectroDynamics::getFields (const vect3d& posv, prec_t t, vect3d& bfield, vect3d& efield, size_t except_electron)
{
  statics.getFields (posv, bfield, efield);

  for ( size_t j = 0; j < delems.size(); j++ )
      delems[j]->addFields (posv, t, bfield, efield);
}

prec_t ElectroDynamics::getPotential (const vect3d& posv, size_t except_particle)
{
  prec_t ret = 0;

  if ( !disable_fields )
    {
      for ( size_t i = 0; i < particles.size(); i++ )
        {
          if ( i == except_particle )
              continue;
          particle& p = particles[i];
          ret += p.charge / (posv - p.pos).length();
        }
      ret *= Kme;
    }

  for ( size_t j = 0; j < delems.size(); j++ )
      delems[j]->addPotential (posv, 0, ret);
  ret += statics.getPotential(posv);

  return ret;
}

#define PARTICLE_CONTACT_RANGE (0.001*0.001)
bool ElectroDynamics::inContact (const vect3d& pos, size_t except)
{
  if ( statics.inContact(pos) )
      return true;

  if ( !disable_fields )
    {
      for ( size_t i = 0; i < particles.size(); i++ )
        {
          if ( i != except && (particles[i].pos - pos).magnitude() <= PARTICLE_CONTACT_RANGE )
              return true;
        }
    }

  return false;
}

prec_t ElectroDynamics::get_elapsed_time ()
{
  return (prec_t)ticks * delta_t;
}

void ElectroDynamics::print ()
{
  prec_t et = get_elapsed_time();
  for ( size_t i = 0; i < particles.size(); i++ )
    {
      vect3d v = particles[i].v;
      vect3d pos = particles[i].pos;
      vect3d vn = v.normal();
      fprintf (printfp, "tick=%lu,time=%.16f,id=%lu,pos=(%f, %f, %f),v=%.3f,vv=(%.2f, %.2f, %.2f)\n",
          (unsigned long)ticks, prec2double(et), (unsigned long)i,
          prec2double(pos.x), prec2double(pos.y), prec2double(pos.z),
          prec2double(v.length()),
          prec2double(vn.x), prec2double(vn.y), prec2double(vn.z));
    }
  fflush (printfp);
}

size_t ElectroDynamics::inject_electron (const vect3d& posv, const vect3d& heading, prec_t electronVolts)
{
  return inject_particle (ELECTRON_MASS, -ELEMENTARY_CHARGE, posv, heading, electronVolts);
}

size_t ElectroDynamics::inject_particle (prec_t mass, prec_t charge, const vect3d& pos,
    const vect3d& heading, prec_t electronVolts)
{
  size_t ret = particles.size();
  particles.push_back (particle());
  particle& p = particles.back();
  p.mass = mass;
  p.charge_by_mass = charge / mass;
  p.charge = charge;
  p.pos = pos;
  if ( electronVolts != 0 )
      p.v = heading.normal() * eV2ms (electronVolts, mass);
  else
      p.v.clear();
  p.adaptlevel = 0;
  getFields (p.pos, p.last_bf, p.last_ef, ret);
  return ret;
}

void ElectroDynamics::reset ()
{
  if ( delems.empty() )
    {
      for ( size_t i = 0; i < size(); i++ )
        {
          particle& p = particles[i];
          getFields (p.pos, p.last_bf, p.last_ef, i);
        }
    }
  else
    {
      for ( size_t i = 0; i < size(); i++ )
        {
          particle& p = particles[i];
          getFields (p.pos, 0, p.last_bf, p.last_ef, i);
        }
    }
}

void ElectroDynamics::update (size_t i, const vect3d& pos, const vect3d& heading, prec_t electronVolts)
{
  particle& p = particles[i];
  p.pos = pos;
  if ( electronVolts != 0 )
      p.v = heading.normal() * eV2ms (electronVolts, p.mass);
  else
      p.v.clear();
}

bool ElectroDynamics::save (const char *path)
{
  FILE *fp = fopen (path, "wb");
  if ( !fp )
      return false;
  FPWriter w(fp);
  w.write ("dYn0", 4);
  w.write (ticks);
  w.write (delta_t);
  w.write (adaptf);
  w.write (madaptf);
  w.write ((size_t)(disable_fields ? 1 : 0));
  w.write (size());
  for ( size_t i = 0; i < size(); i++ )
    {
      particle& p = particles[i];
      w.write (p.pos);
      w.write (p.v);
      w.write (p.mass);
      w.write (p.charge);
    }
  return true;
}

bool ElectroDynamics::load (const char *path)
{
  FILE *fp = fopen (path, "rb");
  if ( !fp )
      return false;
  FPReader r(fp);

  char hdr[4];
  r.read (hdr, 4);
  if ( memcmp (hdr, "dYn0", 4) != 0 )
      return false;
  r.read (ticks);
  r.read (delta_t);
  r.read (adaptf);
  r.read (madaptf);
  r.read (disable_fields);
  size_t psize;
  r.read (psize);
  if ( r.had_errors() )
      return false;
  particles.clear();
  particles.reserve (psize);
  for ( size_t i = 0; i < psize; i++ )
    {
      particles.push_back (particle());
      particle& p = particles.back();
      r.read (p.pos);
      r.read (p.v);
      r.read (p.mass);
      r.read (p.charge);
      p.charge_by_mass = p.charge / p.mass;
      p.adaptlevel = 0;
    }

  update_pmap ();

  reset();

  return !r.had_errors();
}

void ElectroDynamics::addDynamicElement (DynamicElement *elem)
{
  delems.push_back (elem);
}

void ElectroDynamics::erase (size_t i)
{
  particles.erase (particles.begin() + i);
}

void ElectroDynamics::internal_step(size_t i)
{
  particle& p = particles[i];
  vect3d cur_pos = p.pos;
  vect3d cur_v = p.v;
  pp map = pmap[p.adaptlevel];
  size_t count = map.count;
  bool adapt = true;
  bool grow = true;
  bool shrink = true;
  bool first = true;

  while (count)
    {
      vect3d k1 = p.charge_by_mass * (p.last_ef + cur_v % p.last_bf);
      vect3d bf, ef;
      getFields (cur_pos + map.hdt * cur_v + map.dtdt_by_8 * k1, bf, ef, i);
      vect3d k2 = p.charge_by_mass * (ef + (cur_v + map.hdt * k1) % bf);
      if ( adapt )
        {
          prec_t factor = (k1 - k2).magnitude()/k1.magnitude();
          if ( grow && factor > adaptf && p.adaptlevel != ADAPTIVE_STEPS - 1 )
            {
              if ( first )
                {
                  // in case of lower resolution the initial field value will probably diverge too much
                  getFields (cur_pos, p.last_bf, p.last_ef, i);
                  first = false;
                }
              p.adaptlevel++;
              map = pmap[p.adaptlevel];
              count = map.count;
              shrink = false;
              continue;
            }
          else if ( shrink && factor < madaptf && p.adaptlevel != 0 )
            {
              p.adaptlevel--;
              map = pmap[p.adaptlevel];
              count = map.count;
              grow = false;
              continue;
            }
          else
            {
              adapt = false;
            }
        }
      vect3d k3 = p.charge_by_mass * (ef + (cur_v + map.hdt * k2) % bf);
      getFields (cur_pos + map.dt * cur_v + map.dtdt_by_2 * k3, p.last_bf, p.last_ef, i);
      vect3d k4 = p.charge_by_mass * (p.last_ef + (cur_v + map.dt * k3) % p.last_bf);

      cur_pos = cur_pos + map.dt * cur_v + map.dtdt_by_6 * (k1 + k2 + k3);
      cur_v = cur_v + map.dt_by_6 * (k1 + 2 * k2 + 2 * k3 + k4);

      count--;
    }

  if ( disable_fields )
    {
      p.pos = cur_pos;
      p.v = cur_v;
    }
  else
    {
      p.next_pos = cur_pos;
      p.next_v = cur_v;
    }
}

void ElectroDynamics::dynstep(size_t i)
{
  particle& p = particles[i];
  vect3d cur_pos = p.pos;
  vect3d cur_v = p.v;
  pp map = pmap[p.adaptlevel];
  size_t count = map.count;
  bool adapt = true;
  bool grow = true;
  bool shrink = true;
  bool first = true;
  prec_t cur_t = 0;

  while (count)
    {
      vect3d k1 = p.charge_by_mass * (p.last_ef + cur_v % p.last_bf);
      vect3d bf, ef;
      getFields (cur_pos + map.hdt * cur_v + map.dtdt_by_8 * k1, cur_t + map.qhdt, bf, ef, i);
      vect3d k2 = p.charge_by_mass * (ef + (cur_v + map.hdt * k1) % bf);
      if ( adapt )
        {
          prec_t factor = (k1 - k2).magnitude()/k1.magnitude();
          if ( grow && factor > adaptf && p.adaptlevel != ADAPTIVE_STEPS - 1 )
            {
              if ( first )
                {
                  // in case of lower resolution the initial field value will probably diverge too much
                  getFields (cur_pos, p.last_bf, p.last_ef, i);
                  first = false;
                }
              p.adaptlevel++;
              map = pmap[p.adaptlevel];
              count = map.count;
              shrink = false;
              continue;
            }
          else if ( shrink && factor < madaptf && p.adaptlevel != 0 )
            {
              p.adaptlevel--;
              map = pmap[p.adaptlevel];
              count = map.count;
              grow = false;
              continue;
            }
          else
            {
              adapt = false;
            }
        }
      cur_t += map.qdt;
      vect3d k3 = p.charge_by_mass * (ef + (cur_v + map.hdt * k2) % bf);
      getFields (cur_pos + map.dt * cur_v + map.dtdt_by_2 * k3, cur_t, p.last_bf, p.last_ef, i);
      vect3d k4 = p.charge_by_mass * (p.last_ef + (cur_v + map.dt * k3) % p.last_bf);

      cur_pos = cur_pos + map.dt * cur_v + map.dtdt_by_6 * (k1 + k2 + k3);
      cur_v = cur_v + map.dt_by_6 * (k1 + 2 * k2 + 2 * k3 + k4);

      count--;
    }

  if ( disable_fields )
    {
      p.pos = cur_pos;
      p.v = cur_v;
    }
  else
    {
      p.next_pos = cur_pos;
      p.next_v = cur_v;
    }
}

void ElectroDynamics::step (size_t count)
{
  assert (!disable_fields && count == 1 || disable_fields);
  ticks += count;
#ifdef HAVE_THREADS
  if ( particles.size() > 64 && tmgr ) // context switching overhead
    {
      for ( size_t i = 0; i < tasks.size(); i++ )
        {
          tasks[i].c = count;
          tmgr->addTask (&tasks[i]);
        }
      while ( tmgr->run() != 0 );
    }
  else
#endif
  if ( delems.empty() )
    {
      for ( size_t i = 0; i < particles.size(); i++ )
        {
          for ( size_t j = 0; j < count; j++ )
              internal_step(i);
        }
    }
  else
    {
      for ( size_t i = 0; i < particles.size(); i++ )
        {
          for ( size_t j = 0; j < count; j++ )
              dynstep(i);
        }
    }

  if ( !disable_fields )
    {
      // now apply the real change
      for ( size_t i = 0; i < particles.size(); i++ )
        {
          particle& p = particles[i];
          p.pos = p.next_pos;
          p.v = p.next_v;
        }
    }
}

#ifdef HAVE_THREADS
void ElectroDynamics::run_steptask (size_t n, size_t c)
{
  size_t pc = particles.size() / THREADS_CPUCOUNT;
  size_t si = n * pc;
  size_t ei = si + pc;
  if ( n == tasks.size() - 1 )
      ei = particles.size();

  if ( delems.empty() )
    {
      for ( size_t i = si; i < ei; i++ )
        {
          for ( size_t j = 0; j < c; j++ )
              internal_step(i);
        }
    }
  else
    {
      for ( size_t i = si; i < ei; i++ )
        {
          for ( size_t j = 0; j < c; j++ )
              dynstep(i);
        }
    }
}
#endif

prec_t ElectroDynamics::eV2ms (prec_t ev, prec_t mass)
{
  return prec_t_sqrt (2 * ELEMENTARY_CHARGE * ev / mass);
}

prec_t ElectroDynamics::ms2eV (prec_t ms, prec_t mass)
{
  return mass * ms * ms / (2 * ELEMENTARY_CHARGE);
}

size_t ElectroDynamics::ticks_for_distance (prec_t speed, prec_t distance)
{
  prec_t t = distance / speed;
  return (size_t)prec2double(t / delta_t);
}

DynamicElement::~DynamicElement ()
{
}

Dynamic3::Dynamic3 ()
{
  index = 2;
  next();
}

Dynamic3::~Dynamic3 ()
{
}

void Dynamic3::next ()
{
  index++;
  if ( index == 3 )
      index = 0;
  s0 = &get(0);
  s1 = &get(1);
  s2 = &get(2);
}

#define X0 -2
#define X1 -1
#define X2 0

static inline prec_t LGRNG3 (prec_t Y0, prec_t Y1, prec_t Y2, prec_t x)
{
  return Y0 * (x - X1) * (x - X2) * 0.5 - Y1 * (x - X0) * (x - X2) + Y2 * (x - X0) * (x - X1) * 0.5;
}

void Dynamic3::addFields (const vect3d& posv, prec_t t, vect3d& bfield, vect3d& efield)
{
  vect3d bf2, bf1, bf0, ef2, ef1, ef0;
  s0->getFields (posv, bf0, ef0);
  s1->getFields (posv, bf1, ef1);
  s2->getFields (posv, bf2, ef2);

  // printf ("addFields: t=%g, ef0.x=%g, ef1.x=%g, ef2.x=%g, res=%g\n",
  //    t, ef0.x, ef1.x, ef2.x, LGRNG3(ef0.x, ef1.x, ef2.x, t));

  bfield += vect3d(LGRNG3(bf0.x, bf1.x, bf2.x, t), LGRNG3(bf0.y, bf1.y, bf2.y, t), LGRNG3(bf0.z, bf1.z, bf2.z, t));
  efield += vect3d(LGRNG3(ef0.x, ef1.x, ef2.x, t), LGRNG3(ef0.y, ef1.y, ef2.y, t), LGRNG3(ef0.z, ef1.z, ef2.z, t));
}

void Dynamic3::addPotential (const vect3d& posv, prec_t t, prec_t& potential)
{
  prec_t pot0 = s0->getPotential (posv);
  prec_t pot1 = s1->getPotential (posv);
  prec_t pot2 = s2->getPotential (posv);

  potential += LGRNG3 (pot0, pot1, pot2, t);
}


